import cv2
import torch
import requests
import numpy as np
import torchvision
from PIL import Image
from torch import no_grad
import matplotlib.pyplot as plt
from torchvision import transforms


"""
get_predictions函数用于处理和过滤对象检测模型做出的预测。
它接受三个参数：
  1、pred作为模型的原始输出.
  2、阈值用于过滤低置信度的预测.
  3、对象作为可选列表用于过滤基于特定类的预测。
  首先, 该函数通过创建包含类名、检测概率和边界框坐标的元组列表, 将原始预测转换为更可读的格式。然后过滤掉不符合概率阈值的预测, 如果提供了objects参数，
  则只保留与指定对象名称匹配的预测。结果是一个经过过滤的元组列表，可供进一步分析或可视化
"""

# Function to get predictions with optional filtering by object and threshold
def get_predictions(pred, threshold=0.8, objects=None):
  """
   Assign a string name to predicted classes and filter out predictions below a give thresold.Args:
      pred: List containing tuples with class labels, probabilities, and bounding boxes
      thresold: Minimum probability required to consider a prediction valid.
      objects: Optional list of object names to filter predictions.
  Returns:
      List of tuples containing class name, probability, and bounding box for each valid prediction. 
  """
  predicted_classes = [(COCO_INSTANCE_CATEGORY_NAMES[i], p, [(box[0],box[1]),(box[2],box[3])])
                       for i, p, box in zip(list(pred[0]['labels'].numpy()),
                                                 pred[0]['scores'].detach().numpy(),
                                                 list(pred[0]['boxes'].detach().numpy()))]
  predicted_classes = [stuff for stuff in predicted_classes if stuff[1] > threshold]

  if objects and predicted_classes:
    predicted_classes = [(name, p, box) for name, p, box in predicted_classes if name in objects]
  return predicted_classes

# Function to draw bounding boxes around detected objects
def draw_box(predicted_classes, image, rect_th= 1, text_size=1, text_th=1):
  """
  Draw bounding boxes and labels around detected objects in an image.
  Args:
    predicted_classes: List of tuples containing class name, probability, and bounding box.
    image: Image tenser on which boxes and labels will be drawn.
    rect_th: Thickness of the rectangle.
    text_size: Font size of the label text
    text_th: Thickness of the label text.
  """